from textual.app import App, ComposeResult
from textual.widgets import Header, Footer, Input, Button, Static
from textual.containers import VerticalScroll, Horizontal
from textual.reactive import reactive
from textual.message import Message
from modelscope import AutoModelForCausalLM, AutoTokenizer
from flask import Flask, request, Response, redirect


model_name = "Qwen/Qwen2.5-0.5B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
    model_name, torch_dtype="auto", device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)


messages = [
    {
        "role": "system",
        "content": "你是qwen助手",
    }
]


def query_llm(query):
    messages.append({"role": "user", "content": query})
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
    generated_ids = model.generate(**model_inputs, max_new_tokens=512)
    generated_ids = [
        output_ids[len(input_ids) :]
        for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]
    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    messages.append({"role": "assistant", "content": response})
    return response


class ChatApp(App):
    # CSS_PATH = "chat.css"

    messages = reactive([])

    def compose(self) -> ComposeResult:
        yield VerticalScroll(id="message-container")
        
        # 输入框
        with Horizontal(id="horizontal-container"):
            input_t = Input(placeholder="输入...", id="input")
            input_t.styles.width = '80%'
            yield input_t
            
            send_button = Button("发送", id="send-button")
            send_button.styles.max_width = 6
            yield send_button
        
        # yield Footer()

    def on_mount(self) -> None:
        self.query_one("#input").focus()
        self.query_one("#horizontal-container").styles.height=4

    def on_input_submitted(self, event: Input.Submitted) -> None:
        self.send_message(event.value)
        event.input.value = ""

    def on_button_pressed(self, event: Button.Pressed) -> None:
        if event.button.id == "send-button":
            input_widget = self.query_one("#input", Input)
            self.send_message(input_widget.value)
            input_widget.value = ""
            input_widget.focus()

    def send_message(self, message: str) -> None:
        if message.strip():
            self.messages.append(message)

            message_container = self.query_one("#message-container", VerticalScroll)

            # 添加询问
            message_widget = Static(message)
            message_widget.styles.text_align = "right"
            message_container.mount(message_widget)

            # llm回答
            message_widget = Static(query_llm(message))
            message_widget.styles.text_align = "left"
            message_container.mount(message_widget)


class SendMessage(Message):
    """Message sent when a new chat message is added."""

    def __init__(self, message: str) -> None:
        super().__init__()
        self.message = message


if __name__ == "__main__":
    app = ChatApp()
    app.run()
